# Copyright (c) Phigent Robotics. All rights reserved.
import torch
import torch.nn.functional as F
from mmcv.runner import force_fp32

from mmdet3d.ops.bev_pool_v2.bev_pool import TRTBEVPoolv2
from mmdet.models import DETECTORS
from .. import builder
from .bevdet import BEVDet


@DETECTORS.register_module()
class BEVDet_Map(BEVDet):
    r"""BEVDet paradigm for multi-camera 3D object detection.

    Please refer to the `paper <https://arxiv.org/abs/2112.11790>`_

    Args:
        img_view_transformer (dict): Configuration dict of view transformer.
        img_bev_encoder_backbone (dict): Configuration dict of the BEV encoder
            backbone.
        img_bev_encoder_neck (dict): Configuration dict of the BEV encoder neck.
    """

    def __init__(self, 
        map_head=None, 
        **kwargs):
        super(BEVDet_Map, self).__init__(
            **kwargs)

        train_cfg, test_cfg = kwargs['train_cfg'], kwargs['test_cfg']
        if map_head:
            pts_train_cfg = train_cfg.maps if train_cfg else None
            map_head.update(train_cfg=pts_train_cfg)
            pts_test_cfg = test_cfg.maps if test_cfg else None
            map_head.update(test_cfg=pts_test_cfg)
            self.map_head = builder.build_head(map_head)
    
    @staticmethod
    def update_task_weight(loss: dict, w: float):
        for k,v in loss.items():
            v *= w
        return loss

    def forward_train(self,
                      points=None,
                      img_metas=None,
                      gt_bboxes_3d=None,
                      gt_labels_3d=None,
                      gt_labels=None,
                      gt_bboxes=None,
                      img_inputs=None,
                      proposals=None,
                      gt_bboxes_ignore=None,
                      **kwargs):
        """Forward training function.

        Args:
            points (list[torch.Tensor], optional): Points of each sample.
                Defaults to None.
            img_metas (list[dict], optional): Meta information of each sample.
                Defaults to None.
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`], optional):
                Ground truth 3D boxes. Defaults to None.
            gt_labels_3d (list[torch.Tensor], optional): Ground truth labels
                of 3D boxes. Defaults to None.
            gt_labels (list[torch.Tensor], optional): Ground truth labels
                of 2D boxes in images. Defaults to None.
            gt_bboxes (list[torch.Tensor], optional): Ground truth 2D boxes in
                images. Defaults to None.
            img (torch.Tensor optional): Images of each sample with shape
                (N, C, H, W). Defaults to None.
            proposals ([list[torch.Tensor], optional): Predicted proposals
                used for training Fast RCNN. Defaults to None.
            gt_bboxes_ignore (list[torch.Tensor], optional): Ground truth
                2D boxes in images to be ignored. Defaults to None.

        Returns:
            dict: Losses of different branches.
        """
        img_feats, pts_feats, _ = self.extract_feat(
            points, img=img_inputs, img_metas=img_metas, **kwargs)
        losses = dict()
        if hasattr(self, "pts_bbox_head"):
            losses_pts = self.forward_pts_train(img_feats, gt_bboxes_3d,
                                                gt_labels_3d, img_metas,
                                                gt_bboxes_ignore)
            pts_w = self.train_cfg.pts.get("loss_weight", 1.0)
            self.update_task_weight(losses_pts, pts_w)
            losses.update(losses_pts)
        # print(img_feats[0].mean(),losses.keys())
        if hasattr(self,"map_head"):
            loss_map = self.map_head.forward_pts_train(img_feats, 
                                                        kwargs["gt_maps_3d"],
                                                        img_metas)
                                        
            map_w = self.train_cfg.maps.get("loss_weight", 1.0)
            self.update_task_weight(loss_map, map_w)

            losses.update(loss_map)
        # print(img_feats[0].mean(), loss_map.keys())
        return losses

    def forward_test(self,
                     points=None,
                     img_metas=None,
                     img_inputs=None,
                     **kwargs):
        """
        Args:
            points (list[torch.Tensor]): the outer list indicates test-time
                augmentations and inner torch.Tensor should have a shape NxC,
                which contains all points in the batch.
            img_metas (list[list[dict]]): the outer list indicates test-time
                augs (multiscale, flip, etc.) and the inner list indicates
                images in a batch
            img (list[torch.Tensor], optional): the outer
                list indicates test-time augmentations and inner
                torch.Tensor should have a shape NxCxHxW, which contains
                all images in the batch. Defaults to None.
        """
        for var, name in [(img_inputs, 'img_inputs'),
                          (img_metas, 'img_metas')]:
            if not isinstance(var, list):
                raise TypeError('{} must be a list, but got {}'.format(
                    name, type(var)))

        num_augs = len(img_inputs)
        if num_augs != len(img_metas):
            raise ValueError(
                'num of augmentations ({}) != num of image meta ({})'.format(
                    len(img_inputs), len(img_metas)))

        if not isinstance(img_inputs[0][0], list):
            img_inputs = [img_inputs] if img_inputs is None else img_inputs
            points = [points] if points is None else points
            return self.simple_test(points[0], img_metas[0], img_inputs[0],
                                        **kwargs)
        else:
            return self.aug_test(None, img_metas[0], img_inputs[0], **kwargs)

    def aug_test(self, points, img_metas, img=None, rescale=False):
        """Test function without augmentaiton."""
        assert False

    def simple_test(self,
                    points,
                    img_metas,
                    img=None,
                    rescale=False,
                    **kwargs):
        """Test function without augmentaiton."""
        if 0:
            x, rots, trans, intrins, post_rots, post_trans = img[:6]
            x[0].cpu().numpy().tofile('./tmp/img_input.bin')

            pose1 = torch.stack([torch.eye(4) for i in range(post_rots.shape[1])], 0)
            pose2 = torch.stack([torch.eye(4) for i in range(post_rots.shape[1])], 0)
            pose1[:,:3,:3] = post_rots.squeeze(0)
            pose1[...,:3,3] = post_trans.squeeze(0)
            torch.inverse(pose1).cpu().numpy().tofile('./tmp/pose1.bin')

            combine = rots.matmul(torch.inverse(intrins))
            pose2[...,:3,:3] = combine.squeeze(0)
            pose2[...,:3,3] = trans.squeeze(0)
            pose2.cpu().numpy().tofile('./tmp/pose2.bin')

        img_feats, _, _ = self.extract_feat(
            points, img=img, img_metas=img_metas, **kwargs)
        bbox_list = [dict() for _ in range(len(img_metas))]
        losses = dict()
        if hasattr(self, "pts_bbox_head"):
            bbox_pts = self.simple_test_pts(img_feats, 
                                            img_metas,
                                            rescale=rescale)
            for result_dict, pts_bbox in zip(bbox_list, bbox_pts):
                # pts_bbox['boxes_3d'] = pts_bbox['boxes_3d'].tensor #todofix for infer
                result_dict['pts_bbox'] = pts_bbox
                # result_dict['timestamp'] = kwargs['timestamp']
        # print(img_feats[0].mean(),losses.keys())
        if hasattr(self,"map_head"):
            maps_bev = self.map_head.simple_test_maps(img_feats,
                                                     img_metas,
                                                     rescale=rescale)
            for result_dict, maps in zip(bbox_list, maps_bev):
                result_dict['maps'] = maps
                result_dict["points"] = points  
                result_dict.update(kwargs)   
                # result_dict['gt_maps_3d'] = kwargs['gt_maps_3d']      
        return bbox_list

    def forward_dummy(self,img=None):
        """Test function without augmentaiton."""
        img_feats, _, _ = self.extract_feat_dummy(img=img)
        out = []
        if hasattr(self, "pts_bbox_head"):
            outs_box = self.pts_bbox_head(img_feats)
            out.append(outs_box)
        if hasattr(self,"map_head"):
            outs_maps = self.map_head(img_feats)
            out.append(outs_maps)
        return out

    def forward(self, *args, **kwargs):
        if len(args) == 0:
            return super().forward(**kwargs)
        else:
            return self.forward_dummy(img=args[0])

@DETECTORS.register_module()
class BEVDepth_Map(BEVDet_Map):

    def forward_train(self,
                      points=None,
                      img_metas=None,
                      gt_bboxes_3d=None,
                      gt_labels_3d=None,
                      gt_labels=None,
                      gt_bboxes=None,
                      img_inputs=None,
                      proposals=None,
                      gt_bboxes_ignore=None,
                      **kwargs):
        """Forward training function.

        Args:
            points (list[torch.Tensor], optional): Points of each sample.
                Defaults to None.
            img_metas (list[dict], optional): Meta information of each sample.
                Defaults to None.
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`], optional):
                Ground truth 3D boxes. Defaults to None.
            gt_labels_3d (list[torch.Tensor], optional): Ground truth labels
                of 3D boxes. Defaults to None.
            gt_labels (list[torch.Tensor], optional): Ground truth labels
                of 2D boxes in images. Defaults to None.
            gt_bboxes (list[torch.Tensor], optional): Ground truth 2D boxes in
                images. Defaults to None.
            img (torch.Tensor optional): Images of each sample with shape
                (N, C, H, W). Defaults to None.
            proposals ([list[torch.Tensor], optional): Predicted proposals
                used for training Fast RCNN. Defaults to None.
            gt_bboxes_ignore (list[torch.Tensor], optional): Ground truth
                2D boxes in images to be ignored. Defaults to None.

        Returns:
            dict: Losses of different branches.
        """
        img_feats, pts_feats, depth = self.extract_feat(
            points, img=img_inputs, img_metas=img_metas, **kwargs)
        losses = dict()
        gt_depth = kwargs['gt_depth']
        loss_depth = self.img_view_transformer.get_depth_loss(gt_depth, depth)
        losses.update(loss_depth)
        if hasattr(self, "pt_bbox_head"):
            losses_pts = self.forward_pts_train(img_feats, gt_bboxes_3d,
                                                gt_labels_3d, img_metas,
                                                gt_bboxes_ignore)
            pts_w = self.train_cfg.pts.get("loss_weight", 1.0)
            self.update_task_weight(losses_pts, pts_w)
            losses.update(losses_pts)
        if hasattr(self, "map_head"):
            losses_maps = self.map_head.forward_pts_train(img_feats,
                                                         kwargs["gt_maps_3d"],
                                                         img_metas)
            map_w = self.train_cfg.maps.get("loss_weight", 1.0)
            self.update_task_weight(losses_maps, map_w)
            losses.update(losses_maps)

        return losses